Social Distancing for coronavirus

There was a great Washington Post article on social distancing. It comes with simulations to show you the effect of different measures, and I think it's fantastic #SciCom. I want to talk it over with my kids, but I want to be ready for them to ask questions like "what if we did (some cool other idea) for social distancing?" So, before I talk it over with them, I want to have some Python code to replicate it and play with. Here are my initial thoughts:

  • When I tweeted about it, I thought it looked like hard sphere MD. And maybe there are some good stat mech results to be had from that. But I think it looks a bit more like agent based modeling.
  • I want this to be as reproducible as possible, so I'm tempted to use just the scientific python stack. I'm also tempted to use mesa which looks awesome.
  • Ugh. I think this means I write it as an object oriented setup, which makes it a bit less accessible.
  • I'm going to make it a bunch of nested loops, which will be inefficient, but will more easily support different agent types.
  • I'll write one notebook with all of the code, and then another that has it wrapped up as a package.
  • I'll also wrap it up in streamlit because OMFG
    • Note to self: if I do use streamlit, it looks like it might be smart to just stuff everything into a pandas dataframe. Here's an example of wrapping up a matplotlib animation in streamlit.
  • Since physics, model the non-movers as just very high mass
  • NOTE TO SELF: run is fast enough for like 100-1000 particles. The animation is slow. That means I want to pregenerate the trajectory, and then graph that later.

Here's the python setup, using anaconda

conda create --name covid python=3 jupyter jupyterlab numpy scipy matplotlib pandas
conda activate covid
In [1]:
import numpy as np, scipy as sp, pandas as pd
from matplotlib import pyplot as plt
from sklearn.metrics.pairwise import euclidean_distances
from matplotlib import animation
from IPython.display import HTML
from itertools import combinations
from collections import namedtuple
In [2]:
#%load_ext line_profiler
In [3]:
"""
Units: time is measured in 12 hour chunks, but there's no strict correspondence to units here.
"""

class EffectiveArea:
    """Tells us the boundaries of the area that people can live in.
    """
    def __init__(self,):
        self.xmin, self.xmax = 0, 20
        self.ymin, self.ymax = 0, 20

class Person:
    def __init__(self,effectivearea,state='succeptible',distancing=False):
        self.r = 0.15
        self.days_infected = 0
        self._typical_recovery_days = 14
        self.distancing = distancing
        if self.distancing:
            self.m = 1000
        else:
            self.m = 1
        self.state = state
        self.ea = effectivearea
        self.x, self.y = np.random.uniform(self.ea.xmin,self.ea.xmax), np.random.uniform(self.ea.ymin,self.ea.ymax)
        if self.distancing:
            self.vx,self.vy = 0,0
        else:
            self.vx, self.vy = np.random.normal(size=2) # Maxwell Boltzmann?
    def move(self,dt,ea):
        self.x, self.y = self.x + self.vx*dt, self.y + self.vy*dt
        if self.state == 'infected':
            if self.days_infected < self._typical_recovery_days:
                self.days_infected += dt
            else:
                if np.random.uniform() < 0.5:
                    self.state = 'recovered'
                    self.days_infected = 0
                else:
                    self.days_infected += dt
    
def collide(p1,p2):
    if p1.state == 'infected' and p2.state == 'succeptible':
        p2.state = 'infected'
    elif p2.state == 'infected' and p1.state == 'succeptible':
        p1.state = 'infected'

    m1, m2 = p1.m, p2.m
    r1, r2 = np.array([p1.x,p1.y]), np.array([p2.x,p2.y])
    v1, v2 = np.array([p1.vx,p1.vy]), np.array([p2.vx,p2.vy])
    M = m1 + m2
    d = np.linalg.norm(r1 - r2)**2
    u1 = v1 - 2*m2 / M * np.dot(v1-v2, r1-r2) / d * (r1 - r2)
    u2 = v2 - 2*m1 / M * np.dot(v2-v1, r2-r1) / d * (r2 - r1)
    p1.vx,p1.vy = u1
    p2.vx,p2.vy = u2
    
    
    
class Universe:
    def __init__(self,npeople,initial_infection_chance=0.1,
                distancing=0.0):
        self.npeople = npeople
        self.ea = EffectiveArea()
        self.dt = 0.1
        self.data = None # gets set in self.run
        def _state():
            if np.random.uniform() < initial_infection_chance:
                return 'infected'
            return 'succeptible'
        def _distancing():
            if np.random.uniform() < distancing:
                return True
            return False
        self.people = [Person(self.ea,_state(),_distancing()) for i in range(self.npeople)]
        self.color = {'succeptible':0.5,'infected':0.0,'recovered':0.7}
        self.color = {'succeptible':'lightblue','infected':'red','recovered':'green'}
        
    def _step(self):
        points = list(zip([p.x for p in self.people],[p.y for p in self.people]))
        dists = euclidean_distances(points,points)
        close = dists < 2*self.people[0].r
        close = close.tolist()
        for (i,j) in combinations(range(self.npeople),2):
            if close[i][j]: # a bit faster than numpy indexing once things get big.
                collide(self.people[i],self.people[j])
        for p in self.people:
            p.move(self.dt,self.ea)
            if p.x < self.ea.xmin or p.x > self.ea.xmax: 
                p.vx = -p.vx
            if p.y < self.ea.ymin or p.y > self.ea.ymax:
                p.vy = -p.vy
    
    def run(self,steps):
        #x_coords[frame,particle_number]
        x_coords = np.zeros((steps,len(self.people)))
        y_coords = np.zeros((steps,len(self.people)))
        state = np.zeros((steps,len(self.people)),dtype='object')
        
        s,i,r = np.zeros(steps),np.zeros(steps),np.zeros(steps)
        def pop_count(people):
            s,i,r = 0,0,0
            for p in people:
                if p.state == 'succeptible':
                    s += 1
                elif p.state == 'infected':
                    i += 1
                elif p.state == 'recovered':
                    r += 1
            return s,i,r
        s[0],i[0],r[0] = pop_count(self.people)
        x_coords[0] = [p.x for p in self.people]
        y_coords[0] = [p.y for p in self.people]
        state[0] = [p.state for p in self.people]
        
        for step in range(1,steps):
            self._step()
            s[step],i[step],r[step] = pop_count(self.people)
            x_coords[step] = [p.x for p in self.people]
            y_coords[step] = [p.y for p in self.people]
            state[step] = [p.state for p in self.people]
        dtype = namedtuple('RunData',['x','y','state','s','i','r','steps'])
        self.data = dtype(x_coords,y_coords,state,s,i,r,steps)
    def draw(self,ax=None):
        if ax is None:
            fig,ax = plt.subplots(figsize=(5,5))
        scat = ax.scatter([p.x for p in self.people],[p.y for p in self.people],
                   c = [self.color[p.state] for p in self.people],
                   marker='.')
        return scat,
In [4]:
def getanim(u):
    fig,ax = plt.subplots(figsize=(4,4))

    left, width = 0.1, 0.85
    bottom, height = 0.1, 0.65
    spacing = 0.02

    rect_universe = [left, bottom, width, height]
    rect_trend = [left, bottom + height + spacing, width, 0.2]

    ax_universe = plt.axes(rect_universe)
    ax_universe.tick_params(axis='x', which='both',bottom=False,top=False,labelbottom=False)
    ax_universe.axis('off')
    ax_trend = plt.axes(rect_trend)

    s,i,r = u.data.s,u.data.i,u.data.r

    ax_trend.stackplot(range(len(s)), i, s, r, labels=['sick','healthy','recovered'],colors=['red','lightblue','green'])

    scat, = u.draw(ax_universe)
    
    def drawframe(i):
        data = np.column_stack(([u.data.x[i],u.data.y[i]]))
        scat.set_offsets(data)
        colors = np.array([u.color[_] for _ in u.data.state[i]])
        scat.set_color(colors)

        _s,_i,_r = np.zeros(len(u.data.s)),np.zeros(len(u.data.s)),np.zeros(len(u.data.s))
        _s[:i],_i[:i],_r[:i] = u.data.s[:i],u.data.i[:i],u.data.r[:i]
        #ax_trend.stackplot(range(len(_s)), _s, _i, _r, labels=['s','i','r'],colors=['blue','green','yellow'])
        #ax_trend.legend(loc='upper left')

        return scat,


    anim = animation.FuncAnimation(fig, drawframe, frames=u.data.steps,
                                  interval=20, blit=False, repeat=False)
    return anim

Now let's try it

For all of these, we'll do 200 people, 200 steps, initial infection chance 10%

In [5]:
npeople, nsteps = 200, 200
initial_infection_chance = 0.1

No distancing

In [6]:
%%capture
distancing = 0.0
u = Universe(npeople,distancing=distancing,initial_infection_chance=initial_infection_chance)
u.run(nsteps)
anim = getanim(u)
In [7]:
HTML(anim.to_jshtml())
Out[7]:

25% of people distance

In [8]:
%%capture
distancing = 0.25
u = Universe(npeople,distancing=distancing,initial_infection_chance=initial_infection_chance)
u.run(nsteps)
anim = getanim(u)
In [9]:
HTML(anim.to_jshtml())
Out[9]:

50% of people distance

In [10]:
%%capture
distancing = 0.5
u = Universe(npeople,distancing=distancing,initial_infection_chance=initial_infection_chance)
u.run(nsteps)
anim = getanim(u)
In [11]:
HTML(anim.to_jshtml())
Out[11]:

75% of people distance

In [12]:
%%capture
distancing = 0.75
u = Universe(npeople,distancing=distancing,initial_infection_chance=initial_infection_chance)
u.run(nsteps)
anim = getanim(u)
In [13]:
HTML(anim.to_jshtml())
Out[13]:

90% of people distance

In [14]:
%%capture
distancing = 0.90
u = Universe(npeople,distancing=distancing,initial_infection_chance=initial_infection_chance)
u.run(nsteps)
anim = getanim(u)
In [15]:
HTML(anim.to_jshtml())
Out[15]:
In [ ]: